#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import uuid
import glob
from matplotlib import cm
from pathlib import Path
from IPython import display
import glob
from PIL import Image
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import Rectangle
import numpy as np
import math
from .waymo_training_math_tools import *
import Dataloader.src.waymo_transfer_training_to_tf as transfer_tool
from .waymo_transfer_training_to_tf import *
def create_figure_and_axes(dpi, size_pixels, size_inch):

    fig, ax = plt.subplots(1, 1, num=uuid.uuid4())
    # Sets output image to pixel resolution
    #size_inches = size_pixels / dpi
    fig.set_size_inches([size_inch, size_inch])
    fig.set_dpi(dpi)
    fig.set_facecolor('white')
    ax.set_facecolor('white')
    ax.xaxis.label.set_color('black')
    ax.tick_params(axis='x', colors='black')
    ax.yaxis.label.set_color('black')
    ax.tick_params(axis='y', colors='black')
    fig.set_tight_layout(True)
    ax.grid(False)
    return fig, ax
    

def _generate_fig_canvas_image(fig):
  # Returns a [H, W, 3] uint8 np.array image from fig.canvas.tostring_rgb()
  # Just enough margin in the figure to display xticks and yticks
  fig.subplots_adjust(left=0.08, bottom=0.08, right=0.98, top=0.98, wspace=0.0, hspace=0.0)
  fig.canvas.draw()
  data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
  fig.clear()
  return data.reshape(fig.canvas.get_width_height()[::-1] + (3,))


def _generate_best_fit_view_point(scenario):    
    
    #print("max x in scenario: ", np.max(scenario.all_agents_x_list))
    #print("max y in scenario: ", np.max(scenario.all_agents_y_list))
    #print("min x in scenario: ", np.min(scenario.all_agents_x_list))
    #print("min y in scenario: ", np.min(scenario.all_agents_y_list))    
    #print("range x in scenario: ", np.ptp(scenario.all_agents_x_list))
    #print("range y in scenario: ", np.ptp(scenario.all_agents_y_list))
    
    center_x = (np.max(scenario.all_agents_x_list) + np.min(scenario.all_agents_x_list)) / 2
    center_y = (np.max(scenario.all_agents_y_list) + np.min(scenario.all_agents_y_list)) / 2       
    range_x = np.ptp(scenario.all_agents_x_list) + 30
    range_y = np.ptp(scenario.all_agents_y_list) + 30   
    width = max(range_x, range_y)    
    return center_x, center_y, width  


def generate_single_step_plot(dataloader, scenario, time_index, start_time ,end_time, dpi, size_pixels, size_inch, plot_path, tra_list = None, pre_id_list = None):
    
    fig, ax = create_figure_and_axes(dpi, size_pixels, size_inch)
    
    center_x, center_y, width = _generate_best_fit_view_point(scenario)
    
    
    # ax.scatter(lc.lc_features_list[0].C_point_x_list,lc.lc_features_list[0].C_point_y_list, color = 'w', marker = '.', linewidths = 0.2)
    # ========== part1: plot map points ==========       
    #plot map points - white
    ax.scatter(scenario.broken_single_white_list[0], scenario.broken_single_white_list[1], color = 'orange',marker = '.', linewidths = 0.3)
    ax.scatter(scenario.solid_single_white_list[0],  scenario.solid_single_white_list[1],  color = 'orange', marker = '.', linewidths = 0.3)
    ax.scatter(scenario.solid_double_white_list[0],  scenario.solid_double_white_list[1],  color = 'orange', marker = '.', linewidths = 0.3)
    
    white_lane_dict = scenario.passing_yellow_lane_dict
    white_lane_id_list = list(white_lane_dict.keys())
    # # xxx = white_lane_dict[44][0]
    # # yyy = white_lane_dict[44][1]
    # lb,_ = dataloader.get_lane_boundary_byID(130)
    # lbb = lb[41]
    # # #ax.scatter(xxx,  yyy,  color = 'pink',alpha=0.7, s=35,zorder = 20)
    # ax.scatter(lbb[0],  lbb[1],  color = 'pink',alpha=0.7, s=35,zorder = 20)
    # # ax.scatter(rbb[0],  rbb[1],  color = 'pink',alpha=0.7, s=35,zorder = 20)
    # # ax.scatter(tbb[0],  tbb[1],  color = 'pink',alpha=0.7, s=35,zorder = 20)
    for wl in white_lane_id_list:
        if white_lane_dict[wl][0] and white_lane_dict[wl][1]:
            xp = white_lane_dict[wl][0][0]
            yp = white_lane_dict[wl][1][0]
            ax.text(xp,yp,"{:d}".format(int(wl)), fontsize = 8,color='red',zorder = 10)
    
    # plot map points - yellow    
    ax.scatter(scenario.broken_single_yellow_list[0],  scenario.broken_single_yellow_list[1], color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.broken_double_yellow_list[0],  scenario.broken_double_yellow_list[1], color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.solid_single_yellow_list[0],   scenario.solid_single_yellow_list[1],  color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.solid_double_yellow_list[0],   scenario.solid_double_yellow_list[1],  color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.passing_double_yellow_list[0], scenario.passing_double_yellow_list[1],color = 'yellow', marker = '.', linewidths = 0.3)  

    # plot map points - lane - black 
    #print("center lane point number:", len(scenario.lane_black_list[0]))
    if len(scenario.lane_black_list) > 0:
        
        for key in scenario.center_lane_id_dict:
            #print((scenario.center_lane_id_dict[key][0][0], scenario.center_lane_id_dict[key][1][0],int(key)))
            ax.scatter(scenario.center_lane_id_dict[key][0], scenario.center_lane_id_dict[key][1], color = 'black', marker = '.', linewidths = 0.2,zorder=1)
            ax.text(scenario.center_lane_id_dict[key][0][0], scenario.center_lane_id_dict[key][1][0],  "{:d}".format(int(key)), fontsize = 8,color='green',zorder = 2)         

    # for lane in scenario.center_lane_list:
    #     if lane.lane_id == 140:
    #         ax.scatter(lane.point_x_list,lane.point_y_list, color = 'black', marker = '.', linewidths = 0.2) 
    # plot stop sign - red
    if len(scenario.stop_sign_list) > 0:
        ax.scatter(scenario.stop_sign_list[0], scenario.stop_sign_list[1], color = 'red', marker = '^', linewidths = 2)
    
    # plot cross walk - darkorange   
    if len(scenario.cross_walk_list) > 0:
        for crosswalk in scenario.cross_walk_list:
            plt.plot(crosswalk.point_x_list, crosswalk.point_y_list, color = 'darkorange', linestyle = '--', linewidth = 2)  
            plt.plot([crosswalk.point_x_list[0], crosswalk.point_x_list[-1]], [crosswalk.point_y_list[0],crosswalk.point_y_list[-1]], color = 'darkorange', linestyle = '--', linewidth = 2)  
        
    # plot road edge - brown
    if len(scenario.road_edge_list) > 0:
        ax.scatter(scenario.road_edge_list[0], scenario.road_edge_list[1], color = 'brown', marker = '.', linewidths = 0.3)      
      
    # ========== part2: plot agents states ========== 
    
    if pre_id_list is not None:
        if time_index < 76:
            time_end = time_index + 15
        else:
            time_end = 90
        for pre_id in pre_id_list:
            pos_x = dataloader. get_agent_center_x_list_byId(pre_id,time_index,time_end)
            pos_x = pos_x[pos_x != 0.0].tolist()
            pos_y = dataloader. get_agent_center_y_list_byId(pre_id,time_index,time_end)
            pos_y = pos_y[pos_y != 0.0].tolist()
            #ax.scatter(pos_x,pos_y,color = 'red', marker = '.', linewidths = 0.3, zorder = 4)
            ax.plot(pos_x, pos_y, linestyle = '-', color = 'pink',zorder=5)
    # plot single trajectory (a list as input)
    if isinstance(tra_list,list):
        tra_xlist = tra_list[0]
        tra_ylist = tra_list[1]
        ax.scatter(tra_xlist, tra_ylist, color = 'green', marker = '.', linewidths = 0.6,zorder=6)
    
    # plot multiple trajectory (a dict as input)
    if isinstance(tra_list,dict):
        lane_list = list(tra_list.keys())
        for i,lane_id in enumerate(lane_list):
            color = plt.cm.jet(i / len(lane_list))
            tra_x = tra_list[lane_id][0]
            tra_y = tra_list[lane_id][1]
            ax.scatter(tra_x,tra_y,color = color,maker = '.',linewidths = 0.6, zorder = 7)
            ax.text(tra_x[-1], tra_y[-1],lane_id,fontsize = 6,color='red')
            
    agents_to_be_pred = dataloader.get_agents_to_pre()
    
    for agent in scenario.agents_list:
        
        if agent.states_array.size == 0:
            continue    
        
        state = agent.states_array[time_index, :]
            
        if state[13] == '-1':
            continue
        
        agent_x = state[3]
        agent_y = state[4]
        agent_yaw    = radian_to_degree(state[9])
        agent_length = state[6]
        agent_width  = state[7]
        agent_height = state[8]
        agent_velocity  = state[12]
        agent_type = state[2]
        agent_id = state[1]
                   
        # draw surrounding vehicles
        if agent_type == 1 and (agent_id not in agents_to_be_pred):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'blue',
                        facecolor = 'blue',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                       agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)    
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
            # draw surrounding pedestrians        
        if agent_type == 2 and (agent_id not in agents_to_be_pred):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'green',
                        facecolor = 'green',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                    agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            ax.text(x_bottom_left,y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')  
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
        if agent_type == 2 and (agent_id in agents_to_be_pred):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'pink',
                        facecolor = 'pink',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                    agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            ax.text(x_bottom_left,y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')  
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
        if agent_type == 1 and (agent_id  in agents_to_be_pred):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'pink',
                        facecolor = 'pink',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                        agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)    
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
              
        # draw surrounding cyclists        
        if agent_type == 3 and (agent_id not in agents_to_be_pred): 
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'orange',
                        facecolor = 'orange',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                       agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')  
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
        if agent_type == 3 and (agent_id in agents_to_be_pred): 
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'pink',
                        facecolor = 'pink',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                       agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')  
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)  
        # draw self-driving car itself        
        if agent_type == -1: 
            # print("this is a sdc,id is !")
            # print(agent_id)
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            #ax.scatter(agent_x, agent_y, color = 'blue', marker = '.', linewidths = 1.5)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'red',
                        facecolor = 'red',
                        fill=True,
                        lw=1))   
            x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                       agent_length, agent_width)
            ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple')  
            ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
            v1_id,v2_id = find_front_back_vehicle(scenario, time_index, agent_id)
            if v1_id is not None:
                v1 = scenario.agents_list[Id_to_Index(scenario.agents_list,v1_id)]
                v1_x = v1.states_array[time_index,:][3]
                v1_y = v1.states_array[time_index,:][4]
                ax.scatter(v1_x, v1_y, color = 'yellow', marker = '*', linewidths = 1.5, zorder=3)
            if v2_id is not None:
                v2 = scenario.agents_list[Id_to_Index(scenario.agents_list,v2_id)]
                v2_x = v2.states_array[time_index,:][3]
                v2_y = v2.states_array[time_index,:][4]
                ax.scatter(v2_x, v2_y, color = 'pink', marker = '*', linewidths = 1.5, zorder=3)
            '''  
            vehicle1, vehicle2 = find_preceeding_following_in_lane(scenario, 
                                                    time_index,
                                                    agent_id,
                                                    )
            if vehicle1 is not None:
                print("Have found!") 
                x_bottom_left, y_bottom_left = get_bottem_left_point_of_rectangle(vehicle1.states_array[time_index, :][3], 
                                        vehicle1.states_array[time_index, :][4], vehicle1.states_array[time_index, :][6], vehicle1.states_array[time_index, :][7], vehicle1.states_array[time_index, :][9])
                ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        edgecolor = 'blue',
                        facecolor = 'blue',
                        fill=True,
                        lw=1))
            else:
                print("not found!")
            # if vehicle2 is not None: 
            #     print("hjdjdjsdj")
            #     vehicle2_plot_x = [agent_x, vehicle2.states_array[time_index][3]]
            #     vehicle2_plot_y = [agent_y, vehicle2.states_array[time_index][4]] 
 
                ax.plot(vehicle2_plot_x,  vehicle2_plot_y, color = 'red',  linestyle = '--', linewidth = 2) 
            '''
            
            '''
            left_vehicle1, left_vehicle2, right_vehicle1, right_vehicle2 = \
            find_closet_vehicles_in_left_right_lane(scenario, 
                                                    time_index,
                                                    agent_id,
                                                    agent_x, 
                                                    agent_y, 
                                                    agent_length, 
                                                    agent_width, 
                                                    agent_yaw)
            
            left_vehicle1_plot_x = [agent_x, left_vehicle1[0]]
            left_vehicle1_plot_y = [agent_y, left_vehicle1[1]]
            left_vehicle2_plot_x = [agent_x, left_vehicle2[0]]
            left_vehicle2_plot_y = [agent_y, left_vehicle2[1]] 
            
            right_vehicle1_plot_x = [agent_x, right_vehicle1[0]]
            right_vehicle1_plot_y = [agent_y, right_vehicle1[1]]
            right_vehicle2_plot_x = [agent_x, right_vehicle2[0]]
            right_vehicle2_plot_y = [agent_y, right_vehicle2[1]]             
            
            ax.plot(left_vehicle1_plot_x,  left_vehicle1_plot_y, color = 'red',  linestyle = '--', linewidth = 2) 
            ax.plot(left_vehicle2_plot_x,  left_vehicle2_plot_y, color = 'red',  linestyle = '--', linewidth = 2) 
            ax.plot(right_vehicle1_plot_x, right_vehicle1_plot_y, color = 'red', linestyle = '--', linewidth = 2) 
            ax.plot(right_vehicle2_plot_x, right_vehicle2_plot_y, color = 'red', linestyle = '--', linewidth = 2) 
            
        '''    
    # Title.
    title = "agents states"
    ax.set_title(title, fontsize = 20)
    ax.set_xlabel("x(m)", fontsize = 20)
    ax.set_ylabel("y(m)", fontsize = 20)
    ax.tick_params(axis = 'x', labelsize = 20)
    ax.tick_params(axis = 'y', labelsize = 20)
    
    # Set axes.
    size = max(40, width * 1.0)
    ax.axis([-size / 2 + center_x, size / 2 + center_x, -size / 2 + center_y, size / 2 + center_y])
    ax.set_aspect('equal')
    
    image = _generate_fig_canvas_image(fig)
    fig.clear()
    plt.clf() 
    plt.close("all") 
    return image        


    
def generate_all_agents_plot_in_one_scenario(scenario, dpi, size_pixels, size_inch, plot_path):
    
    fig, ax = create_figure_and_axes(dpi, size_pixels, size_inch)
    
    center_x, center_y, width = _generate_best_fit_view_point(scenario)
    
    # ========== part1: plot map points ==========       
    # plot map points - white
    ax.scatter(scenario.broken_single_white_list[0], scenario.broken_single_white_list[1], color = 'w', marker = '.', linewidths = 0.3)
    ax.scatter(scenario.solid_single_white_list[0],  scenario.solid_single_white_list[1],  color = 'w', marker = '.', linewidths = 0.3)
    ax.scatter(scenario.solid_double_white_list[0],  scenario.solid_double_white_list[1],  color = 'w', marker = '.', linewidths = 0.3)     
    
    # plot map points - yellow    
    ax.scatter(scenario.broken_single_yellow_list[0],  scenario.broken_single_yellow_list[1], color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.broken_double_yellow_list[0],  scenario.broken_double_yellow_list[1], color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.solid_single_yellow_list[0],   scenario.solid_single_yellow_list[1],  color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.solid_double_yellow_list[0],   scenario.solid_double_yellow_list[1],  color = 'yellow', marker = '.', linewidths = 0.3)  
    ax.scatter(scenario.passing_double_yellow_list[0], scenario.passing_double_yellow_list[1],color = 'yellow', marker = '.', linewidths = 0.3)  

    #plot map points - lane - black 
    #print("center lane point number:", len(scenario.lane_black_list[0]))
    
    ax.scatter(scenario.lane_black_list[0], scenario.lane_black_list[1], color = 'red', marker = '.', linewidths = 0.3)      
    # for lane in scenario.center_lane_list:
    #     if lane.lane_id == 140:
    #         ax.scatter(lane.point_x_list,lane.point_y_list, color = 'black', marker = '.', linewidths = 0.2) 
    # plot stop sign - red
    ax.scatter(scenario.stop_sign_list[0], scenario.stop_sign_list[1], color = 'red', marker = '^', linewidths = 10)
    
    # plot cross walk - darkorange    
    ax.scatter(scenario.cross_walk_list[0], scenario.cross_walk_list[1], color = 'darkorange', marker = '.', linewidths = 3)  
    
    # plot road edge - brown
    ax.scatter(scenario.road_edge_list[0], scenario.road_edge_list[1], color = 'brown', marker = '.', linewidths = 0.3)      
      
    # ========== part2: plot agents states ==========        
    for agent in scenario.agents_list:
        states = agent.states_array
        states_mask = agent.states_mask
        masked_x = states[:, 3][states_mask]
        masked_y = states[:, 4][states_mask]
                
        for state in states:
            if state.valid == '-1':
                continue
            agent_x = state.x
            agent_y = state.y
            agent_yaw    = radian_to_degree(state.heading)
            agent_length = state.length
            agent_width  = state.width
            agent_height = state.height
            agent_velocity  = math.sqrt(state.velocity_x**2 + state.velocity_y**2)
            agent_type = state.agent_type
            agent_id = state.agent_id
            
                 
            # draw surrounding vehicles
            if agent_type == 1:
              x_bottom_left, y_bottom_left = get_bottem_left_point_of_rectangle(agent_x, 
                                          agent_y, agent_length, agent_width, agent_yaw)
              ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                          angle = agent_yaw,
                          edgecolor = 'blue',
                          facecolor = 'blue',
                          fill=True,
                          lw=1))
             
              x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, 
                                                         y_bottom_left, agent_length, agent_width)
              ax.arrow(x, y, delta_x, delta_y, width = 0.1)
              ax.text(x + delta_x, y + delta_y, "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
              
            # draw surrounding pedestrians        
            if agent_type == 2:
              x_bottom_left, y_bottom_left = get_bottem_left_point_of_rectangle(agent_x, 
                                              agent_y, agent_length, agent_width, agent_yaw)
              ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                          angle = agent_yaw,
                          edgecolor = 'green',
                          facecolor = 'green',
                          fill=True,
                          lw=1))
              x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                         agent_length, agent_width)
              ax.arrow(x, y, delta_x, delta_y, width = 0.1)     
              ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
              
            # draw surrounding cyclists        
            if agent_type == 3: 
              x_bottom_left, y_bottom_left = get_bottem_left_point_of_rectangle(agent_x, 
                                              agent_y, agent_length, agent_width, agent_yaw)
              ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                          angle = agent_yaw,
                          edgecolor = 'orange',
                          facecolor = 'orange',
                          fill=True,
                          lw=1))     
              x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                         agent_length, agent_width)
              ax.arrow(x, y, delta_x, delta_y, width = 0.1)     
              ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
              
            # draw self-driving car itself        
            if agent_type == -1: 
              #print("there is a sdc !")
              x_bottom_left, y_bottom_left = get_bottem_left_point_of_rectangle(agent_x, 
                                              agent_y, agent_length, agent_width, agent_yaw)
              ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                          angle = agent_yaw,
                          edgecolor = 'red',
                          facecolor = 'red',
                          fill=True,
                          lw=1))   
              x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                         agent_length, agent_width)
              ax.arrow(x, y, delta_x, delta_y, width = 0.1)    
              ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
            index += 1
        
    ax.set_xlabel("x(m)", fontsize = 20)
    ax.set_ylabel("y(m)", fontsize = 20)
    ax.tick_params(axis = 'x', labelsize = 20)
    ax.tick_params(axis = 'y', labelsize = 20)
    
    # Set axes.
    size = max(40, width * 1.0)
    ax.axis([-size / 2 + center_x, size / 2 + center_x, -size / 2 + center_y, size / 2 + center_y])
    ax.set_aspect('equal')
    
    image = _generate_fig_canvas_image(fig) 
    plt.clf()  
    fig.clear()
    plt.close("all")
    return image        


def create_save_single_plot(single_image, plot_path, scenario_id, plot_number, dpi):

    fig, ax = plt.subplots()
    
    #size_inches = 1000 / dpi
    size_inches = 20
    fig.set_size_inches([size_inches, size_inches])

    red_patch = mpatches.Patch(color='red', label='autonomous vehicle')
    blue_patch = mpatches.Patch(color='blue', label='surrounding vehicles')
    green_patch = mpatches.Patch(color='green', label='pedestrians')
    orange_patch = mpatches.Patch(color='orange', label='cyclists')  

    # Title.
    scenario_id = str(scenario_id)
    plot_title = scenario_id + ": agent states at timestep:" + '{:03}'.format(plot_number)
    ax.set_title(plot_title, fontsize = 20)    
    ax.legend(handles=[red_patch, blue_patch, green_patch, orange_patch], fontsize = 14)
    ax.imshow(single_image)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.grid('off')
    plt.tight_layout()
    if plot_number == None:    
        file_name = plot_path + scenario_id
        #file_name = plot_path
        fig.savefig(file_name, dpi = dpi)
        fig.clear()
        plt.clf()  
        plt.close("all")      
    else:    
        file_name = plot_path + scenario_id + "_" + '{:03}'.format(plot_number)
        #file_name = plot_path
        fig.savefig(file_name, dpi = dpi)   
        fig.clear()
        plt.clf()  
        plt.close("all")
    #plt.show()


def visualization( data_path, dataloader,  tra_result_list = None, pre_id_list = None, start_time = 21, end_time = 90, generate_individual_images = True, generate_animation = True):
    waymo_dataset_list = []
    waymo_dataset_list.append(data_path)
    data_set_index = 7
    for data_set in waymo_dataset_list:
        if generate_individual_images:
            print("waymo open dataset analysis begin!")
            size_pixels = 300
            dpi = 400
            size_inch = 22
                   
            scenarios_to_plot_list = [0]
            scenarios_summary= transfer_tool.process_waymo_training_dataset(data_set, scenarios_to_plot_list)
            time_index_list = scenarios_summary.scenarios_list[0].timestep_list
            states_index_list = time_index_list   
            scenario_list = scenarios_summary.scenarios_list        
            
            scenario_index = 0       
            for scenario in scenario_list:                     
    
                plot_path = "/home/mdmp/ppp/waymo_training_plot/" + "dataset_" + str(data_set_index) + "/scenario_" + str(scenario_index) + "/"
                isExist = os.path.exists(plot_path)               
                if isExist:
                    scenario_index += 1
                    continue
                if not isExist:             
                  os.makedirs(plot_path)
                  print("The new directory is created!")
                
                plot_number = 0
                for time_index in time_index_list:
                    
                    single_step_plot_of_scenario = generate_single_step_plot(dataloader, scenarios_summary.scenarios_list[0], time_index_list[time_index], start_time, end_time, dpi, size_pixels, size_inch, plot_path, tra_result_list, pre_id_list)
                    create_save_single_plot(single_step_plot_of_scenario, plot_path, scenarios_summary.scenarios_list[0].scenario_id, states_index_list[time_index], dpi)
                    print("the " + str(time_index) + "th figure saved")
                    print("plot number is :", plot_number)
                    plot_number += 1
                    #if plot_number > 10:
                        #break           
                    
                if generate_animation:   
                    
                    single_images_path = plot_path + "*.png"
                    gif_output_path = "/home/mdmp/ppp/waymo_training_plot/" + "dataset_" + str(data_set_index) + "/scenario_" + str(scenario_index) + ".gif"
                    
                    single_images = (Image.open(image_file) for image_file in sorted(glob.glob(single_images_path)))
                    first_image = next(single_images)  # extract first image from iterator
                    first_image.save(fp = gif_output_path, 
                                     format = 'GIF', 
                                     append_images = single_images,
                                     save_all = True, 
                                     duration = 50, 
                                     loop = 0)
                    first_image = None
                scenario_index += 1
            data_set_index += 1
            
            
def visualize_on_step(data_path,dataloader,tra_result_list = None,time_step = 0,plot_path = '/home/mdmp/ppp/waymo_plot/'):
    print("waymo open dataset test begin!")
    scenario_index = 0
    size_pixels = 300
    dpi = 400
    size_inch = 22
    start_time = 0
    end_time = 90
    pre_id_list = None
    scenarios_to_plot_list = [0]
    scenarios_summary= transfer_tool.process_waymo_training_dataset(data_path, scenarios_to_plot_list)
    
    isExist = os.path.exists(plot_path)               
    if not isExist:             
        os.makedirs(plot_path)
        print("The new directory is created!")
    single_step_plot_of_scenario = generate_single_step_plot(dataloader, scenarios_summary.scenarios_list[0], time_step, start_time, end_time, dpi, size_pixels, size_inch, plot_path, tra_result_list, pre_id_list)
                
    create_save_single_plot(single_step_plot_of_scenario, plot_path, scenarios_summary.scenarios_list[0].scenario_id, time_step, dpi)
    
    print("the " + str(time_step) + "th figure saved")

    
    return True




def visualize_one_step_muti_trajectory(data_path,dataloader,tra_result_dict = None,time_step = 0):
    print("waymo open dataset test begin!")
    scenario_index = 0
    size_pixels = 300
    dpi = 400
    size_inch = 22
    start_time = 0
    end_time = 90
    pre_id_list = None
    scenarios_to_plot_list = [0]
    scenarios_summary= transfer_tool.process_waymo_training_dataset(data_path, scenarios_to_plot_list)
    plot_path = "/home/mdmp/ppp/waymo_plot/"
    isExist = os.path.exists(plot_path)               
    if not isExist:             
        os.makedirs(plot_path)
        print("The new directory is created!")
    single_step_plot_of_scenario = generate_single_step_plot(dataloader, scenarios_summary.scenarios_list[0], time_step, start_time, end_time, dpi, size_pixels, size_inch, plot_path, tra_result_dict, pre_id_list)
                
    create_save_single_plot(single_step_plot_of_scenario, plot_path, scenarios_summary.scenarios_list[0].scenario_id, time_step, dpi)
    
    print("the " + str(time_step) + "th figure saved")

    
    return True




def trajectory_visualization(data_path,dataloader,trajectory_dict,interval_num, save_path = '/home/mdmp/ppp/waymo_trajectory_vis',confidence_threshold = 0.1,start_time = 20, end_time = 90,is_show_gt = True):
    print("waymo open dataset test begin!")
    scenario_index = 0
    size_pixels = 300
    dpi = 400
    size_inch = 22
    pre_id_list = None
    scenarios_to_plot_list = [0]
    scenarios_summary= transfer_tool.process_waymo_training_dataset(data_path, scenarios_to_plot_list)
    # plot_path = "/home/mdmp/ppp/waymo_trajectory_vis/"
    # isExist = os.path.exists(plot_path)
    # if not isExist:
    #     os.makedirs(plot_path)
    #     print("The new directory is created!")
        
    scenario = scenarios_summary.scenarios_list[0]
    fig, ax = create_figure_and_axes(dpi, size_pixels, size_inch)
    
    center_x, center_y, width = _generate_best_fit_view_point(scenarios_summary.scenarios_list[0])
    
    good_grey_rgb = (142/255, 140/255, 132/255)
    
     #plot map points - white
    ax.scatter(scenario.broken_single_white_list[0], scenario.broken_single_white_list[1], color = good_grey_rgb, alpha = 0.2,s = 5 )
    ax.scatter(scenario.solid_single_white_list[0],  scenario.solid_single_white_list[1],  color = good_grey_rgb, alpha = 0.2,s = 5 )
    ax.scatter(scenario.solid_double_white_list[0],  scenario.solid_double_white_list[1],color = good_grey_rgb, alpha = 0.2,s = 5 )
    
    # plot map points - yellow    
    ax.scatter(scenario.broken_single_yellow_list[0],  scenario.broken_single_yellow_list[1], color = good_grey_rgb, alpha = 0.2,s = 5 )  
    ax.scatter(scenario.broken_double_yellow_list[0],  scenario.broken_double_yellow_list[1], color = good_grey_rgb, alpha = 0.2,s = 5 )  
    ax.scatter(scenario.solid_single_yellow_list[0],   scenario.solid_single_yellow_list[1], color = good_grey_rgb, alpha = 0.2,s = 5 )  
    ax.scatter(scenario.solid_double_yellow_list[0],   scenario.solid_double_yellow_list[1],  color = good_grey_rgb, alpha = 0.2,s = 5 )  
    ax.scatter(scenario.passing_double_yellow_list[0], scenario.passing_double_yellow_list[1],color = good_grey_rgb, alpha = 0.2,s = 5 )  
    
    # plot_short_line_list = scenario.plot_short_line_list
    # for lane_plt in plot_short_line_list:
    #     if lane_plt:
    #         ax.plot(lane_plt[0], lane_plt[1], color = good_grey_rgb, alpha = 0.2,linewidth = 2.0)
    # ax.plot(scenario.solid_single_white_list[0],  scenario.solid_single_white_list[1],  color = good_grey_rgb, alpha = 0.2,linewidth = 2.0 )
    # ax.plot(scenario.solid_double_white_list[0],  scenario.solid_double_white_list[1],color = good_grey_rgb, alpha = 0.2,linewidth = 2.0 )
    
    # # plot map points - yellow    
    # ax.plot(scenario.broken_single_yellow_list[0],  scenario.broken_single_yellow_list[1], color = good_grey_rgb, alpha = 0.2,linewidth = 2.0 )  
    # ax.plot(scenario.broken_double_yellow_list[0],  scenario.broken_double_yellow_list[1], color = good_grey_rgb, alpha = 0.2,linewidth = 2.0 )  
    # ax.plot(scenario.solid_single_yellow_list[0],   scenario.solid_single_yellow_list[1], color = good_grey_rgb, alpha = 0.2,linewidth = 2.0)  
    # ax.plot(scenario.solid_double_yellow_list[0],   scenario.solid_double_yellow_list[1],  color = good_grey_rgb, alpha = 0.2,linewidth = 2.0)
    # ax.plot(scenario.passing_double_yellow_list[0], scenario.passing_double_yellow_list[1],color = good_grey_rgb, alpha = 0.2,linewidth = 2.0)

        
    # plot road edge - brown
    if len(scenario.road_edge_list) > 0:
        ax.scatter(scenario.road_edge_list[0], scenario.road_edge_list[1], color = good_grey_rgb, alpha = 0.2,s = 10 )      
        
    agents_pre_list = list(trajectory_dict.keys())
    
    good_blue_rgb = (51/255,200/255,255/255)
    
    good_green_rgb = (153/255,200/255,100/255)
    
    good_purple_rgb = (164/255,120/255,188/255)
    
    good_red_rgb = (175/255,63/255,22/255)
        
    for agent in scenario.agents_list:
        
        if agent.states_array.size == 0:
            continue    
        
        state = agent.states_array[start_time, :]
            
        if state[13] == '-1':
            continue
        
        agent_x = state[3]
        agent_y = state[4]
        agent_yaw    = radian_to_degree(state[9])
        agent_length = state[6]
        agent_width  = state[7]
        agent_height = state[8]
        agent_velocity  = state[12]
        agent_type = state[2]
        agent_id = state[1]
                   
        # draw surrounding vehicles
        if (agent_type in [-1,1]) and (agent_id in agents_pre_list):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple',alpha = 0.7,zorder = 3)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_green_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))

            if is_show_gt:
                ground_truth_x = dataloader.get_agent_center_x_list_byId(agent_id)[start_time:end_time]
                
                ground_truth_y = dataloader.get_agent_center_y_list_byId(agent_id)[start_time:end_time]

                base_color1 = (0.5,0.5,1.0)
                target_color1 = (0,0,0.5)
                
                g1 = ground_truth_x[::interval_num]
                g2 = ground_truth_y[::interval_num]
                for i, (x, y) in enumerate(zip(g1,g2)):
                    color = tuple(base + (target - base) * (i / len(ground_truth_x)) for base, target in zip(base_color1, target_color1))
                    ax.scatter(x, y, color=color, alpha=0.6, s=15,zorder = 4)

            
            # 初始颜色（浅黄色）
            base_color2 = (1.0, 0.9, 0.0)
            
            # 目标颜色（橙色）
            target_color2 = (1.0, 0.1, 0.0)
            confidence_list = trajectory_dict[int(agent_id)][0]
            tra_list = trajectory_dict[int(agent_id)][1:]
            for confidence, trajectory in zip(confidence_list, tra_list):
                if confidence >= confidence_threshold:
                    tra_x = trajectory[0][::interval_num]
                    tra_y = trajectory[1][::interval_num]
                    for i, (x, y) in enumerate(zip(tra_x,tra_y)):
                        color = tuple(base + (target - base)  * (i / len(tra_x)) for base, target in zip(base_color2, target_color2))
                        ax.scatter(x, y, color=color, alpha=confidence, s=25,zorder = 5)
                    confidence = round(confidence,4)
                    ax.text(tra_x[-1], tra_y[-1],  confidence,fontsize = 12,color='red')
            
                        
            
                 
            
        if (agent_type in [-1,1]) and (not (agent_id in agents_pre_list)):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            # ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='yellow')
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_blue_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))  
            # ground_truth_x = dataloader.get_agent_center_x_list_byId(agent_id)[start_time:end_time]
            
            # ground_truth_y = dataloader.get_agent_center_y_list_byId(agent_id)[start_time:end_time]

            # base_color1 = (0.5,0.5,1.0)
            # target_color1 = (0,0,0.5)
            
            # g1 = ground_truth_x[::interval_num]
            # g2 = ground_truth_y[::interval_num]
            # for i, (x, y) in enumerate(zip(g1,g2)):
            #     color = tuple(base + (target - base) * (i / len(ground_truth_x)) for base, target in zip(base_color1, target_color1))
            #     ax.scatter(x, y, color=color, alpha=0.5, s=15,zorder = 4)
        if agent_type == 2:
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_purple_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))   
            #x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                    #agent_length, agent_width)
            # ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            # ax.text(x_bottom_left,y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='yellow')  
            # ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
                    
        # draw surrounding cyclists        
        if agent_type == 3: 
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_red_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))   
    # Title.
    title = "agents states"
    ax.set_title(title, fontsize = 20)
    ax.set_xlabel("x(m)", fontsize = 20)
    ax.set_ylabel("y(m)", fontsize = 20)
    ax.tick_params(axis = 'x', labelsize = 20)
    ax.tick_params(axis = 'y', labelsize = 20)
    
    # Set axes.
    size = max(40, width * 1.0)
    ax.axis([-size / 2 + center_x, size / 2 + center_x, -size / 2 + center_y, size / 2 + center_y])
    ax.set_aspect('equal')
    
    image = _generate_fig_canvas_image(fig)
    fig.clear()
    plt.clf() 
    plt.close("all") 
    
    
    create_save_single_plot(image, save_path, scenarios_summary.scenarios_list[0].scenario_id, 0, dpi)
    
    print("the figure saved")

    
    return True


def trajectory_visualization_pos(data_path,dataloader,position_dict,interval_num, save_path = '/home/mdmp/ppp/waymo_trajectory_vis',start_time = 10, end_time = 90,is_show_gt = True):
    print("waymo open dataset test begin!")
    scenario_index = 0
    size_pixels = 300
    dpi = 400
    size_inch = 22
    pre_id_list = None
    scenarios_to_plot_list = [0]
    scenarios_summary= transfer_tool.process_waymo_training_dataset(data_path, scenarios_to_plot_list)
    # plot_path = "/home/mdmp/ppp/waymo_trajectory_vis/"
    # isExist = os.path.exists(plot_path)
    # if not isExist:
    #     os.makedirs(plot_path)
    #     print("The new directory is created!")
        
    scenario = scenarios_summary.scenarios_list[0]
    fig, ax = create_figure_and_axes(dpi, size_pixels, size_inch)
    
    center_x, center_y, width = _generate_best_fit_view_point(scenarios_summary.scenarios_list[0])
    
    good_grey_rgb = (142/255, 140/255, 132/255)
    
     #plot map points - white
    ax.scatter(scenario.broken_single_white_list[0], scenario.broken_single_white_list[1], color = good_grey_rgb, alpha = 0.3,s = 20 )
    ax.scatter(scenario.solid_single_white_list[0],  scenario.solid_single_white_list[1],  color = good_grey_rgb, alpha = 0.3,s = 20 )
    ax.scatter(scenario.solid_double_white_list[0],  scenario.solid_double_white_list[1],color = good_grey_rgb, alpha = 0.3,s = 20 )
    
    # plot map points - yellow    
    ax.scatter(scenario.broken_single_yellow_list[0],  scenario.broken_single_yellow_list[1], color = good_grey_rgb, alpha = 0.3,s = 20 )  
    ax.scatter(scenario.broken_double_yellow_list[0],  scenario.broken_double_yellow_list[1], color = good_grey_rgb, alpha = 0.3,s = 20 )  
    ax.scatter(scenario.solid_single_yellow_list[0],   scenario.solid_single_yellow_list[1], color = good_grey_rgb, alpha = 0.3,s = 20 )  
    ax.scatter(scenario.solid_double_yellow_list[0],   scenario.solid_double_yellow_list[1],  color = good_grey_rgb, alpha = 0.3,s = 20 )  
    ax.scatter(scenario.passing_double_yellow_list[0], scenario.passing_double_yellow_list[1],color = good_grey_rgb, alpha = 0.3,s = 20 )  

    # plot map points - lane - black 
    #print("center lane point number:", len(scenario.lane_black_list[0]))
    # if len(scenario.lane_black_list) > 0:
        
    #     for key in scenario.center_lane_id_dict:
    #         #print((scenario.center_lane_id_dict[key][0][0], scenario.center_lane_id_dict[key][1][0],int(key)))
    #         ax.scatter(scenario.center_lane_id_dict[key][0], scenario.center_lane_id_dict[key][1], color = good_grey_rgb, alpha = 0.5,s = 15 ,zorder=1)
    #         ax.text(scenario.center_lane_id_dict[key][0][0], scenario.center_lane_id_dict[key][1][0],  "{:d}".format(int(key)), fontsize = 8,color='green',zorder = 2)         

    # for lane in scenario.center_lane_list:
    #     if lane.lane_id == 140:
    #         ax.scatter(lane.point_x_list,lane.point_y_list, color = 'black', marker = '.', linewidths = 0.2) 
    # plot stop sign - red
    if len(scenario.stop_sign_list) > 0:
        ax.scatter(scenario.stop_sign_list[0], scenario.stop_sign_list[1], color = 'red', marker = '^', linewidths = 2)
    
    # # plot cross walk - darkorange   
    # if len(scenario.cross_walk_list) > 0:
    #     for crosswalk in scenario.cross_walk_list:
    #         plt.plot(crosswalk.point_x_list, crosswalk.point_y_list, color = 'darkorange', linestyle = '--', linewidth = 2)  
    #         plt.plot([crosswalk.point_x_list[0], crosswalk.point_x_list[-1]], [crosswalk.point_y_list[0],crosswalk.point_y_list[-1]], color = 'darkorange', linestyle = '--', linewidth = 2)  
        
    # plot road edge - brown
    if len(scenario.road_edge_list) > 0:
        ax.scatter(scenario.road_edge_list[0], scenario.road_edge_list[1], color = good_grey_rgb, alpha = 0.5,s = 20 )      
        
    agents_pre_list = list(position_dict.keys())
    
    good_blue_rgb = (91/255,145/255,198/255)
    
    good_green_rgb = (87/255,219/255,155/255)
    
    good_purple_rgb = (164/255,120/255,188/255)
    
    good_red_rgb = (175/255,63/255,22/255)
        
    for agent in scenario.agents_list:
        
        if agent.states_array.size == 0:
            continue    
        
        state = agent.states_array[start_time, :]
            
        if state[13] == '-1':
            continue
        
        agent_x = state[3]
        agent_y = state[4]
        agent_yaw    = radian_to_degree(state[9])
        agent_length = state[6]
        agent_width  = state[7]
        agent_height = state[8]
        agent_velocity  = state[12]
        agent_type = state[2]
        agent_id = state[1]
                   
        # draw surrounding vehicles
        if agent_type == 1 and (agent_id in agents_pre_list):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='purple',alpha = 0.6,zorder = 3)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_green_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))

            if is_show_gt:
                if agent_id ==2214:
                    ground_truth_x = dataloader.get_agent_center_x_list_byId(agent_id)[start_time:end_time]
                    
                    ground_truth_y = dataloader.get_agent_center_y_list_byId(agent_id)[start_time:end_time]

                    base_color1 = (0.5,0.5,1.0)
                    target_color1 = (0,0,0.5)
                    
                    g1 = ground_truth_x[::interval_num]
                    g2 = ground_truth_y[::interval_num]
                    for i, (x, y) in enumerate(zip(g1,g2)):
                        color = tuple(base + (target - base) * (i / len(ground_truth_x)) for base, target in zip(base_color1, target_color1))
                        ax.scatter(x, y, color=color, alpha=0.5, s=15,zorder = 4)

            
            # 初始颜色（浅黄色）
            base_color2 = (1.0, 1.0, 0.5)

            # 目标颜色（橙色）
            target_color2 = (1.0, 0.0, 0.0)
            position_list = position_dict[int(agent_id)]
            xlist = position_list[0]
            ylist = position_list[1]
            tra_x = xlist[::interval_num]
            tra_y = ylist[::interval_num]


            for i, (x, y) in enumerate(zip(tra_x,tra_y)):
                color = tuple(base + (target - base) * (i / len(tra_x)) for base, target in zip(base_color2, target_color2))
                ax.scatter(x, y, color=color, alpha=0.7, s=35,zorder = 5)
                # ax.text(tra_x[-1], tra_y[-1],  confidence,fontsize = 6,color='red')
            
            
                 
            
        if agent_type == 1 and (not (agent_id in agents_pre_list)):
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            # ax.text(x_bottom_left, y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='yellow')
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_blue_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))  
            if is_show_gt:
                ground_truth_x = dataloader.get_agent_center_x_list_byId(agent_id)[start_time:end_time]
                
                ground_truth_y = dataloader.get_agent_center_y_list_byId(agent_id)[start_time:end_time]

                base_color1 = (0.5,0.5,1.0)
                target_color1 = (0,0,0.5)
                base_color2 = (1.0, 1.0, 0.5)

            # 目标颜色（橙色）
                target_color2 = (1.0, 0.0, 0.0)
                g1 = ground_truth_x[::interval_num]
                g2 = ground_truth_y[::interval_num]
                for i, (x, y) in enumerate(zip(g1,g2)):
                    color = tuple(base + (target - base) * (i / len(ground_truth_x)) for base, target in zip(base_color1, target_color1))
                    ax.scatter(x, y, color=color, alpha=0.5, s=15,zorder = 4)
        if agent_type == 2:
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_purple_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))   
            #x,y,delta_x,delta_y = calculate_arrow_base(agent_yaw, x_bottom_left, y_bottom_left, 
                                                    #agent_length, agent_width)
            # ax.arrow(x, y, delta_x, delta_y, width = 0.1,zorder=2)  
            # ax.text(x_bottom_left,y_bottom_left,  "{:d}".format(int(agent_id)), fontsize = 6,color='yellow')  
            # ax.text(x + delta_x, y + delta_y,  "{:.2f}".format(agent_velocity) + "m/s", fontsize = 12)
                    
        # draw surrounding cyclists        
        if agent_type == 3: 
            x_bottom_left,y_bottom_left = get_bottem_left_point_of_rectangle_origin(agent_x,agent_y,agent_length,agent_width,agent_yaw)
            ax.add_patch(Rectangle((x_bottom_left, y_bottom_left), agent_length, agent_width, 
                        angle = agent_yaw,
                        facecolor = good_red_rgb,
                        alpha = 0.5,
                        fill=True,
                        lw=1))   

    
    
    
    # Title.
    title = "agents states"
    ax.set_title(title, fontsize = 20)
    ax.set_xlabel("x(m)", fontsize = 20)
    ax.set_ylabel("y(m)", fontsize = 20)
    ax.tick_params(axis = 'x', labelsize = 20)
    ax.tick_params(axis = 'y', labelsize = 20)
    
    # Set axes.
    size = max(40, width * 1.0)
    ax.axis([-size / 2 + center_x, size / 2 + center_x, -size / 2 + center_y, size / 2 + center_y])
    ax.set_aspect('equal')
    
    image = _generate_fig_canvas_image(fig)
    fig.clear()
    plt.clf() 
    plt.close("all") 
    
    
    create_save_single_plot(image, save_path, scenarios_summary.scenarios_list[0].scenario_id, 0, dpi)
    
    print("the figure saved")

    
    return True